#include "VulkanPixelObjects.h"

#include "Render/Widget/VulkanViewportWidget.h"
#include "Render/Mesh/VulkanMesh.h"
#include "Render/PreRender/PreRenderMesh.h"
#include "Render/VulkanMemory/Buffers.h"
#include "Render/VulkanMemory/Images.h"
#include "Render/VulkanMath.h"
#include "VulkanPixelObjectGlobal.h"
#include "PreRenderPixelObject.h"
#include "PreRenderPixelObjectMap.h"
#include "Shaders/CppMacros.h"

#include "Async/PackagedTask.h"
#include "PixelObject/PixelObjectMap.h"


CVulkanPixelObjectGroup::CVulkanPixelObjectGroup()
{
    ViewInfoBuffer=std::make_shared<
        TVulkanUniformBuffer<
            SPixelObjectViewUniformBuffer
        >>();
}

CVulkanPixelObjectGroup::~CVulkanPixelObjectGroup()
{
    //回收framebuffer
    CVulkanContext::Get()->GetRecycleObjects()->Framebuffers.Add(FrameBuffer);
}

void CVulkanPixelObjectGroup::PreProcess(
    std::shared_ptr<CPreRenderPixelObjectGroup> InPixelObjectGroup
    , CVulkanViewportWidget * InViewport)
{

    ViewportSize=InViewport->ViewportSize;
    RenderTargetSize.width=InPixelObjectGroup->MapWidth;
    RenderTargetSize.height=InPixelObjectGroup->MapHeight;

    {
        //创建颜色图像,只有一个通道 uint16
        if(!ColorImage)
        {
            ColorImage=std::make_shared<CVulkanImage>();
        }

        vk::ImageCreateInfo ImageCreateInfo=
            vk::ImageCreateInfo()
            .setImageType(vk::ImageType::e2D)
            .setExtent(vk::Extent3D(
                RenderTargetSize.width
                ,RenderTargetSize.height
                ,1
            ))
            .setMipLevels(1)
            .setArrayLayers(1)
            .setSamples(vk::SampleCountFlagBits::e1)
            .setTiling(vk::ImageTiling::eOptimal)
            .setUsage(vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eTransferSrc)
            .setSharingMode(vk::SharingMode::eExclusive)
            .setInitialLayout(vk::ImageLayout::eUndefined)
            .setFormat(vk::Format::eR16Uint)
            .setQueueFamilyIndexCount(0)
            .setPQueueFamilyIndices(nullptr)
            .setFlags(vk::ImageCreateFlags());

        ColorImage->SetImageParams(
            ImageCreateInfo
        );

        //创建颜色图像视图
        vk::ImageViewCreateInfo ColorImageViewCreateInfo=
            vk::ImageViewCreateInfo()
            .setImage(ColorImage->GetImage())
            .setViewType(vk::ImageViewType::e2D)
            .setFormat(ColorImage->GetFormat())
            .setSubresourceRange(
                vk::ImageSubresourceRange(
                    vk::ImageAspectFlagBits::eColor
                    ,0
                    ,1
                    ,0
                    ,1
                )
            );
        ColorImageView=CVulkanContext::Get()->GetDevice().createImageView(ColorImageViewCreateInfo);


        //创建深度图
        if(!DepthImage)
        {
            DepthImage=std::make_shared<CVulkanImage>();
        }

        DepthImage->SetDepthImageParams(
            RenderTargetSize.width
            ,RenderTargetSize.height
        );

        //创建深度图视图
        vk::ImageViewCreateInfo DepthImageViewCreateInfo=
            vk::ImageViewCreateInfo()
            .setImage(DepthImage->GetImage())
            .setViewType(vk::ImageViewType::e2D)
            .setFormat(DepthImage->GetFormat())
            .setSubresourceRange(
                vk::ImageSubresourceRange(
                    vk::ImageAspectFlagBits::eDepth
                    ,0
                    ,1
                    ,0
                    ,1
                )
            );

        DepthImageView=CVulkanContext::Get()->GetDevice().createImageView(DepthImageViewCreateInfo);
    }
    
    NotifyFormerFrameFinished();
    RenderTargetID=InPixelObjectGroup->PixelObjectMap->ID;
    {
        //创建最终渲染目标数据缓冲区,用于保存渲染结果,将像素数据传给CPU
        auto NewDataSize=RenderTargetSize.width*RenderTargetSize.height*sizeof(SPixelObjectID);
        if(RenderTargetBuffer.GetSize()!=NewDataSize)
        {
            RenderTargetBuffer.ClearBuffer();
        }

        if(!RenderTargetBuffer.Buffer)
        {
            vk::BufferCreateInfo BufferCreateInfo=
                vk::BufferCreateInfo()
                .setSize(NewDataSize)
                .setUsage(vk::BufferUsageFlagBits::eTransferDst)
                .setSharingMode(vk::SharingMode::eExclusive)
                .setQueueFamilyIndexCount(0)
                .setPQueueFamilyIndices(nullptr);
            
            RenderTargetBuffer.Buffer=CVulkanContext::Get()->GetDevice().createBuffer(BufferCreateInfo);
            
            vk::MemoryRequirements MemoryRequirements=
                CVulkanContext::Get()->GetDevice().getBufferMemoryRequirements(RenderTargetBuffer.Buffer);

            auto MemoryTypeIndex=
                CVulkanContext::Get()->GetMemoryTypeIndex(
                    MemoryRequirements
                    ,vk::MemoryPropertyFlagBits::eHostVisible|vk::MemoryPropertyFlagBits::eHostCoherent
                );

            RenderTargetBuffer.MemoryBlock=
                CVulkanContext::Get()->GetMemoryPools()->AllocateBlock(
                    EVulkanMemoryUsage::MappedTexture
                    ,MemoryTypeIndex
                    ,MemoryRequirements);

            CVulkanContext::Get()->GetDevice().bindBufferMemory(
                RenderTargetBuffer.Buffer
                ,RenderTargetBuffer.MemoryBlock.Memory
                ,RenderTargetBuffer.MemoryBlock.Offset
            );
        }

    }


    //创建帧缓冲
    {
        vk::FramebufferCreateInfo FrameBufferCreateInfo;
        FrameBufferCreateInfo.renderPass = CVulkanPixelObjectGlobal::Get()->RenderPass;
        FrameBufferCreateInfo.attachmentCount = 2;
        vk::ImageView Attachments[] = { 
                ColorImageView,DepthImageView};
        FrameBufferCreateInfo.pAttachments = Attachments;
        FrameBufferCreateInfo.width = RenderTargetSize.width;
        FrameBufferCreateInfo.height = RenderTargetSize.height;
        FrameBufferCreateInfo.layers = 1;

        FrameBuffer = CVulkanContext::Get()->GetDevice().createFramebuffer(FrameBufferCreateInfo);
    }

    //更新视角信息
    {
        SPixelObjectViewUniformBuffer ViewInfo;
        (InViewport->ProjectionMatrix*InViewport->ViewMatrix)
            .ToFloatArray(ViewInfo.ProjectViewMatrix);
        ViewInfoBuffer->WriteData(ViewInfo);
    } 

    //统计描述符集大小
    {
        SVulkanDescriptorPoolSize PoolSize;
        PoolSize.SetNum=1;
        PoolSize+=CVulkanPixelObjectGlobal::Get()->ViewDescriptorLayout.DescriptorNums;
        CVulkanContext::Get()->GetAutoResetDescriptorPool()->MeshInstancePool.PoolSize+=PoolSize;
    }


    //创建MeshBatch
    MeshBatchs.Resize(InPixelObjectGroup->Meshs.Num());
    for(int64_t i=0;i<InPixelObjectGroup->Meshs.Num();++i)
    {
        MeshBatchs[i].PreProcess(InPixelObjectGroup->Meshs[i], InViewport);
    }

    //创建ForegroundMeshBatch
    ForegroundMeshBatchs.Resize(InPixelObjectGroup->ForegroundMeshs.Num());
    for (int64_t i = 0; i < InPixelObjectGroup->ForegroundMeshs.Num(); ++i)
    {
        ForegroundMeshBatchs[i].PreProcess(InPixelObjectGroup->ForegroundMeshs[i], InViewport);
    }



}

void CVulkanPixelObjectGroup::Draw(vk::CommandBuffer InCmdBuffer)
{
    CreateViewInfoDescriptorSet();
    for(auto& MeshBatch:MeshBatchs)
    {
        MeshBatch.CreateDescriptorSet();
    }
    for (auto& MeshBatch : ForegroundMeshBatchs)
    {
        MeshBatch.CreateDescriptorSet();
    }

    vk::RenderPassBeginInfo RenderPassBeginInfo;
    RenderPassBeginInfo.renderPass = CVulkanPixelObjectGlobal::Get()->RenderPass;
    RenderPassBeginInfo.framebuffer = FrameBuffer;
    RenderPassBeginInfo.renderArea.offset = vk::Offset2D(0, 0);
    RenderPassBeginInfo.renderArea.extent = RenderTargetSize;
    RenderPassBeginInfo.clearValueCount = 2;
    vk::ClearValue ClearValues[2];
    SPixelObjectID InValidObjectID = std::numeric_limits<SPixelObjectID>::max();
    ClearValues[0].color.float32[0] = InValidObjectID;
    ClearValues[0].color.float32[1] = 0.0f;
    ClearValues[0].color.float32[2] = 0.0f;
    ClearValues[0].color.float32[3] = 0.0f;
    ClearValues[1].depthStencil.depth = 1.0f;
    ClearValues[1].depthStencil.stencil = 0;
    RenderPassBeginInfo.pClearValues = ClearValues;

    InCmdBuffer.beginRenderPass(RenderPassBeginInfo, vk::SubpassContents::eInline);

    //设置视口
    vk::Viewport Viewport;
    Viewport.width = RenderTargetSize.width;
    Viewport.height = RenderTargetSize.height;
    Viewport.minDepth = 0.0f;
    Viewport.maxDepth = 1.0f;
    InCmdBuffer.setViewport(0, Viewport);

    //设置裁剪
    vk::Rect2D Scissor;
    Scissor.extent = RenderTargetSize;
    InCmdBuffer.setScissor(0, Scissor);

    //设置管线
    InCmdBuffer.bindPipeline(vk::PipelineBindPoint::eGraphics
        , CVulkanPixelObjectGlobal::Get()->Pipeline);

    //绘制MeshBatch
    for(auto& MeshBatch:MeshBatchs)
    {
        MeshBatch.Draw(InCmdBuffer,ViewInfoDescriptorSet);
    }

    //绘制ForegroundMeshBatch
    for (auto& MeshBatch : ForegroundMeshBatchs)
    {
        MeshBatch.Draw(InCmdBuffer, ViewInfoDescriptorSet);
    }
      
    InCmdBuffer.endRenderPass();

    //复制colorimage到 RenderTargetBuffer
    {
        vk::ImageMemoryBarrier Barrier;
        Barrier.srcAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
        Barrier.dstAccessMask = vk::AccessFlagBits::eTransferRead;
        Barrier.oldLayout = vk::ImageLayout::eColorAttachmentOptimal;
        Barrier.newLayout = vk::ImageLayout::eTransferSrcOptimal;
        Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        Barrier.image = ColorImage->GetImage();
        Barrier.subresourceRange.aspectMask = vk::ImageAspectFlagBits::eColor;
        Barrier.subresourceRange.baseMipLevel = 0;
        Barrier.subresourceRange.levelCount = 1;
        Barrier.subresourceRange.baseArrayLayer = 0;
        Barrier.subresourceRange.layerCount = 1;

        InCmdBuffer.pipelineBarrier(
            vk::PipelineStageFlagBits::eColorAttachmentOutput
            , vk::PipelineStageFlagBits::eTransfer
            , vk::DependencyFlags()
            , nullptr
            , nullptr
            , Barrier
        );

        vk::ImageSubresourceLayers ImageSubresourceLayers;
        ImageSubresourceLayers.aspectMask = vk::ImageAspectFlagBits::eColor;
        ImageSubresourceLayers.mipLevel = 0;
        ImageSubresourceLayers.baseArrayLayer = 0;
        ImageSubresourceLayers.layerCount = 1;

        vk::BufferImageCopy BufferImageCopy;
        BufferImageCopy.bufferOffset = 0;
        BufferImageCopy.bufferRowLength = 0; //像素值紧密排列
        BufferImageCopy.bufferImageHeight = 0;
        BufferImageCopy.imageSubresource = ImageSubresourceLayers;
        BufferImageCopy.imageOffset = vk::Offset3D(0, 0, 0);
        BufferImageCopy.imageExtent = vk::Extent3D(
            RenderTargetSize.width
            , RenderTargetSize.height
            , 1
        );

        InCmdBuffer.copyImageToBuffer(
            ColorImage->GetImage()
            , vk::ImageLayout::eTransferSrcOptimal
            , RenderTargetBuffer.Buffer
            , BufferImageCopy
        );

        Barrier.srcAccessMask = vk::AccessFlagBits::eTransferRead;
        Barrier.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
        Barrier.oldLayout = vk::ImageLayout::eTransferSrcOptimal;
        Barrier.newLayout = vk::ImageLayout::eColorAttachmentOptimal;

        InCmdBuffer.pipelineBarrier(
            vk::PipelineStageFlagBits::eTransfer
            , vk::PipelineStageFlagBits::eColorAttachmentOutput
            , vk::DependencyFlags()
            , nullptr
            , nullptr
            , Barrier
        );
    }



}

void CVulkanPixelObjectGroup::CreateViewInfoDescriptorSet()
{
    //创建视角描述符集
    {
        vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo=vk::DescriptorSetAllocateInfo()
            .setDescriptorPool(CVulkanContext::Get()->GetAutoResetDescriptorPool()->MeshInstancePool.Pool)
            .setDescriptorSetCount(1)
            .setPSetLayouts(&CVulkanPixelObjectGlobal::Get()->ViewDescriptorLayout.Layout);
        ViewInfoDescriptorSet=CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo)[0];
       
    }
    
    //更新视角描述符集
    {
        vk::DescriptorBufferInfo DescriptorBufferInfo=
            vk::DescriptorBufferInfo()
            .setBuffer(ViewInfoBuffer->GetBuffer())
            .setOffset(0)
            .setRange(ViewInfoBuffer->GetSize());
     
        vk::WriteDescriptorSet WriteDescriptorSet=
            vk::WriteDescriptorSet()
            .setDstSet(ViewInfoDescriptorSet)
            .setDstBinding(0)
            .setDstArrayElement(0)
            .setDescriptorType(vk::DescriptorType::eUniformBuffer)
            .setDescriptorCount(1)
            .setPBufferInfo(&DescriptorBufferInfo);

        CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSet, {});
    }

}

void CVulkanPixelObjectGroup::NotifyFormerFrameFinished()
{
    //std::lock_guard<std::mutex> Lock(Mutex);
    //添加一个主线程任务，通知上一帧已经渲染完毕
    //这个任务会在主线程执行
    if(!RenderTargetBuffer.Buffer
    )
    {
        //上一帧没有渲染，不需要通知
        return;
    }
    assert(RenderTargetID);

    TVector<uint8_t> Pixel;
    Pixel.Resize(RenderTargetBuffer.GetSize());

    //拷贝数据
    memcpy(
        Pixel.Data()
        ,RenderTargetBuffer.MemoryBlock.MappedMemory
        ,RenderTargetBuffer.GetSize());

    

    //创建一个主线程任务，通知上一帧已经渲染完毕
    uint32_t ViewportWidth=ViewportSize.width;
    uint32_t ViewportHeight=ViewportSize.height;
    uint32_t Width=RenderTargetSize.width;
    uint32_t Height=RenderTargetSize.height;
    uint32_t Pitch=RenderTargetBuffer.GetSize()/RenderTargetSize.height;
    TThreadSafeWeakPtr<CPixelObjectMapID> WeakLogicThreadMap=RenderTargetID;
    auto Task=WH::CreateMainThreadTask<void>(
        [
            ViewportWidth
            ,ViewportHeight
            ,Width
            ,Height
            ,Pitch
            ,MovedPixel=std::move(Pixel)
            ,WeakLogicThreadMap
        ](bool&) mutable->void 
        {
            auto LogicThreadObjectWrapper=WeakLogicThreadMap.Lock();
            if(!LogicThreadObjectWrapper)
            {
                //逻辑线程对象的包装器已经销毁
                return;
            }

            auto Map=LogicThreadObjectWrapper->Map.Lock();
            if(!Map)
            {
                //逻辑线程对象已经销毁
                return;
            }

            Map->WriteData(
                ViewportWidth
                ,ViewportHeight
                ,Width
                ,Height
                ,Pitch
                ,std::move(MovedPixel)
            );

        }
    );
    Task->Start();
}

CVulkanPixelObjectMeshBatch::CVulkanPixelObjectMeshBatch()
{
    //创建模型缓冲区
    ModelUniformBuffer= std::make_shared<CVulkanUniformBuffer>();
}

void CVulkanPixelObjectMeshBatch::PreProcess(std::shared_ptr<CPreRenderPixelObjectMesh> InMeshInstance, CVulkanViewportWidget *InViewport)
{
    Mesh = CVulkanContext::Get()->ToVulkanMesh(InMeshInstance->Mesh->WriteInData);

    InstanceNum=InMeshInstance->ModelMatrices.Num();
    BatchNum=CVulkanMath::CalculateBatchNum(InstanceNum,MESH_BATCH_SIZE);


    // ModelMatrices = InMeshInstance->ModelMatrices;
    // ObjectIDs = InMeshInstance->ObjectIDs;

    //更新模型矩阵
    TVector<CVulkanPixelObjectGlobal::SModelUniformBuffer> ModelBufferData;
    ModelBufferData.Resize(InstanceNum);
    for(int32_t i=0;i<InstanceNum;++i)
    {
        InMeshInstance->ModelMatrices[i].ToFloatArray(ModelBufferData[i].ModelMatrix);
        ModelBufferData[i].ObjectID= InMeshInstance->ObjectIDs[i];
    }
    auto DataSize= ModelBufferData.Num()
                 *sizeof(CVulkanPixelObjectGlobal::SModelUniformBuffer);
    ModelUniformBuffer->Resize(DataSize);
    ModelUniformBuffer->WriteData(ModelBufferData.Data(), DataSize);


    CountDescriptorSet();
}

void CVulkanPixelObjectMeshBatch::Draw(
    vk::CommandBuffer InCmdBuffer
    ,vk::DescriptorSet InViewInfoDescriptorSet
    )
{
    for(int32_t BatchIndex=0;BatchIndex<BatchNum;BatchIndex++)
    {
        InCmdBuffer.bindDescriptorSets(vk::PipelineBindPoint::eGraphics
            , CVulkanPixelObjectGlobal::Get()->PipelineLayout
            , 0
            , { ModelBufferDescriptorSets[BatchIndex]
                , InViewInfoDescriptorSet
             }
            , {}
        );
        
        //这个小批次的大小
        int32_t BatchSize=
            CMath::Min(MESH_BATCH_SIZE,InstanceNum-BatchIndex*MESH_BATCH_SIZE);

        Mesh->RecordDrawCommands(InCmdBuffer,BatchSize,BatchIndex*MESH_BATCH_SIZE);
    }
}

void CVulkanPixelObjectMeshBatch::CountDescriptorSet()
{
    SVulkanDescriptorPoolSize PoolSize;
    PoolSize.SetNum=BatchNum;
    PoolSize+=
        CVulkanPixelObjectGlobal::Get()->DescriptorLayout.DescriptorNums* PoolSize.SetNum;
    CVulkanContext::Get()->GetAutoResetDescriptorPool()->MeshInstancePool.PoolSize+=PoolSize;

}

void CVulkanPixelObjectMeshBatch::CreateDescriptorSet()
{
    //创建描述符集
    vk::DescriptorSetAllocateInfo DescriptorSetAllocateInfo;
    {
        DescriptorSetAllocateInfo.descriptorPool = CVulkanContext::Get()->GetAutoResetDescriptorPool()->MeshInstancePool.Pool;
        DescriptorSetAllocateInfo.descriptorSetCount = BatchNum;
        DescriptorSetAllocateInfo.pSetLayouts = &CVulkanPixelObjectGlobal::Get()->DescriptorLayout.Layout;
        
        auto Sets=CVulkanContext::Get()->GetDevice().allocateDescriptorSets(DescriptorSetAllocateInfo);
        ModelBufferDescriptorSets.Resize(BatchNum);
        for(int32_t i=0;i<BatchNum;++i)
        {
            ModelBufferDescriptorSets[i]=Sets[i];
        }
    }

    //更新描述符集
    for(int32_t BatchIndex=0;BatchIndex<BatchNum;++BatchIndex)
    {
        //这个小批次的大小
        int32_t BatchSize=
            CMath::Min(MESH_BATCH_SIZE,InstanceNum-BatchIndex*MESH_BATCH_SIZE);

        vk::DescriptorBufferInfo DescriptorBufferInfo;
        DescriptorBufferInfo.buffer=ModelUniformBuffer->GetBuffer();
        DescriptorBufferInfo.offset=BatchIndex*MESH_BATCH_SIZE*sizeof(CVulkanPixelObjectGlobal::SModelUniformBuffer);
        DescriptorBufferInfo.range=BatchSize*sizeof(CVulkanPixelObjectGlobal::SModelUniformBuffer);

        vk::WriteDescriptorSet WriteDescriptorSet;
        WriteDescriptorSet.dstSet=ModelBufferDescriptorSets[BatchIndex];
        WriteDescriptorSet.dstBinding=0;
        WriteDescriptorSet.dstArrayElement=0;
        WriteDescriptorSet.descriptorType=vk::DescriptorType::eUniformBuffer;
        WriteDescriptorSet.descriptorCount=1;
        WriteDescriptorSet.pBufferInfo=&DescriptorBufferInfo;

        CVulkanContext::Get()->GetDevice().updateDescriptorSets(WriteDescriptorSet, {});
    }
}
